import streamlit as st
import pandas as pd
from create_plan9 import create_plan
from mongo_database9 import update_data

import re
import io


# 设置页面配置
st.set_page_config(layout="wide")

st.subheader("内运排班系统")
st.write("请输入以下信息，用于生成今日排班表")

#初始化任务
if not st.session_state:
    st.session_state = {
        '铝卷' : {'count': None,
                  'discard_drop_count': 1,
                  'discard_count' : 1},
        '铝棒' : {'count': None, 'load_times': []},
        '单顶' : {'count': None, 'load_times': []},
        '双零箔': {
            'count': None,
            'four_point_two': None,
            'ten_point_five': None,
            'thirteen': None,
            'load_time': '7:50'

        },
        '废铝卷':{
            'count': None,
            'scania': None,
            'benz': None
        },
        '长沙甩挂人员': {'name': None},
        '常德甩挂人员': {'name': None},
        '休息人员名单': {'rest_names':[]},
        '休息车辆': {'rest_car_numbers':[]}
    }
st.subheader("铝卷信息")
col1, col2, col3 =  st.columns(3)
with col1 :
    st.session_state['铝卷']['count'] = st.number_input(
        '铝卷数量',min_value = 0, max_value = 15, key = 'aluminum_roll_count'
    )
with col2 :
    st.session_state['铝卷']['discard_drop_count'] = st.number_input(
        '带废铝卷数量(甩废铝)',min_value = 0, max_value = st.session_state['铝卷']['count'], key = 'discard_drop_count'
    )
with col3 :
    st.session_state['铝卷']['discard_count'] = st.number_input(
        '带废铝卷数量(废铝)',min_value = 0, max_value = st.session_state['铝卷']['count'] - st.session_state['铝卷']['discard_drop_count'],key = 'discard_count'
    )

st.subheader("铝棒信息")
aluminum_rod_count =st.number_input(
        '铝棒数量',min_value = 0, max_value = 8, key = 'aluminum_rod_count'
    )

if aluminum_rod_count != st.session_state['铝棒']['count']:
    st.session_state['铝棒']['count'] = aluminum_rod_count
    st.session_state['铝棒']['load_times'] = [None] * aluminum_rod_count
    st.rerun()

# 根据铝棒数量显示时间输入框
if st.session_state['铝棒']['count'] > 0:
    st.subheader("为每趟铝棒输入装货时间")
    # 创建列来组织输入框
    cols = st.columns(2)

for i in range(aluminum_rod_count):
    with cols[i%2]:
        # 使用唯一的key
        time_key = f"aluminum_rod_time_{i}"
        # 获取当前值或使用None
        current_value = st.session_state['铝棒']['load_times'][i] if i < len(
            st.session_state['铝棒']['load_times']) else None

        # 时间输入框
        selected_time = st.time_input(
            f"第 {i + 1} 趟铝棒装货时间",
            value=current_value,
            step = 1800,
            key=time_key
        )
        if selected_time is not None:
        # 使用 strftime 格式化为 HH:MM 字符串
            time_str = selected_time.strftime('%H:%M')
        else:
            time_str = selected_time
        # 更新session state中的时间
        st.session_state['铝棒']['load_times'][i] = time_str

st.subheader("单顶信息")
single_top_count =st.number_input(
        '单顶数量',min_value = 0, max_value = 4, key = 'single_top_count'
    )

if single_top_count != st.session_state['单顶']['count']:
    st.session_state['单顶']['count'] = single_top_count
    st.session_state['单顶']['load_times'] = [None] * single_top_count
    st.rerun()

# 根据单顶数量显示时间输入框
if st.session_state['单顶']['count'] > 0:
    st.subheader("为每趟单顶输入装货时间")
    # 创建列来组织输入框
    cols = st.columns(2)

for i in range(single_top_count):
    with cols[i%2]:
        # 使用唯一的key
        time_key = f"single_top_time_{i}"
        # 获取当前值或使用None
        current_value = st.session_state['单顶']['load_times'][i] if i < len(
            st.session_state['单顶']['load_times']) else None

        # 时间输入框
        selected_time = st.time_input(
            f"第 {i + 1} 趟单顶装货时间",
            value=current_value,
            step = 1800,
            key=time_key
        )
        if selected_time is not None:
        # 使用 strftime 格式化为 HH:MM 字符串
            time_str = selected_time.strftime('%H:%M')
        else:
            time_str = selected_time
        # 更新session state中的时间
        st.session_state['单顶']['load_times'][i] = time_str

st.subheader("双零箔信息")
st.session_state['双零箔']['count'] = st.number_input(
    '双零箔数量', min_value=0, max_value = 6, key='double_zero_count'
)
col1, col2, col3 =  st.columns(3)
with col1 :
    st.session_state['双零箔']['four_point_two'] = st.number_input(
        '4.2米车辆数量', min_value=0, max_value = st.session_state['双零箔']['count'], key='four_point_two'
    )
with col2 :
    st.session_state['双零箔']['ten_point_five'] = st.number_input(
        '10.5米车辆数量',min_value = 0, max_value = 1, key = 'ten_point_five'
    )
with col3 :
    st.session_state['双零箔']['thirteen'] = st.number_input(
        '13米车辆数量',min_value = 0, max_value = 4, key = 'thirteen'
    )


st.subheader("废铝卷信息")
col1, col2, col3 =  st.columns(3)
with col1 :
    st.session_state['废铝卷']['count'] = st.number_input(
        '废铝卷数量',min_value = 0, max_value = 6, key = 'discard_aluminum_roll_count'
    )
with col2 :
    st.session_state['废铝卷']['scania'] = st.number_input(
        '斯堪尼亚车数量',min_value = 0, max_value = st.session_state['废铝卷']['count'], key = 'scania_count'
    )
with col3 :
    st.session_state['废铝卷']['benz'] = st.number_input(
        '奔驰车数量',min_value = 0, max_value = 2,  key = 'benz_count'
    )
st.subheader("甩挂人员信息")

left_col, right_col = st.columns(2)
changsha_list = ['刘锦红','李胜果','张学海','唐基平','徐已兵','卓兵华',
                 '毛天宝','郭磊','王云','郭城','陈贵平','丁冬','刘伟',
                 '向宇晴','燕为胜','王德芹','杨军','陈百胜','曹乐魁']
changde_list = ['邹君','龚亚龙','陈建忠','郭斌','徐爱明']
with left_col :
    st.session_state['长沙甩挂人员']['name'] = st.selectbox(
            '长沙甩挂人员',
            options = changsha_list,
            key = 'changsha_drop_roll_name'
        )
with right_col:
    st.session_state['常德甩挂人员']['name'] = st.selectbox(
            '常德甩挂人员',
            options = changde_list,
            key = 'changde_drop_roll_name'
    )


left_col, right_col = st.columns(2)
with left_col :
    st.subheader('休息人员名单')
    rest_names_input = st.text_input(
            "请输入休息人员姓名，多个姓名用逗号分隔，例如“丁冬，杨军”", key="rest_names"
        )
    if rest_names_input:
        st.session_state['休息人员名单']['rest_names'] = [name.strip() for name in re.split(r'[，,;\s/\\]+', rest_names_input)]
    else:
        st.session_state['休息人员名单']['rest_names'] = []

with right_col :
    st.subheader('休息车辆名单')
    rest_car_numbers_input = st.text_input(
            "请输入休息车辆的车牌号，多个车牌号用逗号分隔，例如“湘AM9210，湘AK1377”", key="rest_car_numbers"
        )
    if rest_car_numbers_input:
        st.session_state['休息车辆']['rest_car_numbers'] = [number.strip() for number in re.split(r'[，,;\s/\\]+', rest_car_numbers_input)]
    else:
        st.session_state['休息车辆']['rest_car_numbers'] = []

if st.button('确认生成计划', type='primary'):
    try:
        result_df,assigned_drivers_df, assigned_vehicles_df, drivers_list, vehicles_list = create_plan(today_tasks=st.session_state)
        for col in result_df.columns:
            if result_df[col].dtype == 'object':
                result_df[col] = result_df[col].astype(str)
        left_col, mid_col, right_col = st.columns(3)
        with left_col:
            stats_data = {
                '货物类型': ['铝卷数量', '铝棒数量', '单顶数量', '双零箔数量'],
                '数量': [
                    result_df['货物类型'].str.contains('铝卷', na=False).sum(),
                    result_df['货物类型'].value_counts().get('铝棒', 0),
                    result_df['货物类型'].value_counts().get('单顶', 0),
                    result_df['货物类型'].value_counts().get('双零箔', 0)
                ]
            }
            stats_df = pd.DataFrame(stats_data)

            st.write('已排班货物数量')
            st.table(stats_df)
        with mid_col:
            st.write('未安排人员名单信息')
            st.write(drivers_list)
        with right_col:
            st.write('未安排车辆名单信息')
            st.write(vehicles_list)
        st.success('计划生成成功！')
        st.session_state['df'] = result_df[[
        '序','司机姓名', '车牌号','发车时间', '装货时间', '货物类型', '始发地','目的地'
,'收车时间'
    ]].copy()
        st.session_state['assigned_drivers'] = result_df.copy()
        st.session_state['assigned_vehicles'] = result_df.copy()

    except Exception as e:
        st.error(f'生成计划失败，请重新输入任务信息: {e}')

# 始终显示数据编辑器（如果数据存在）
if 'df' in st.session_state and not st.session_state['df'].empty:
    # 使用 data_editor 并自动更新 session_state
    st.session_state['df'] = st.data_editor(
        st.session_state['df'],
        num_rows="dynamic",
        width='stretch',
        height= 600,
        hide_index=True,
        key='main_editor'
    )

    # 创建两列布局，用于放置按钮
    col1, col2 = st.columns(2)

    with col1:
        # 添加保存按钮
        if st.button('保存数据表，更新历史数据', use_container_width=True):
            try:
                # 调用更新函数
                update_data(st.session_state['df'], st.session_state['assigned_drivers'],
                            st.session_state['assigned_vehicles'])
                st.success('数据保存成功！')
            except Exception as e:
                st.error(f'数据保存失败: {e}')

    with col2:

        # 创建Excel文件
        output = io.BytesIO()
        with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
            st.session_state['df'].to_excel(writer, index=False, sheet_name='计划数据表')

        # 下载Excel文件
        st.download_button(
            label="下载数据表 (Excel)",
            data=output.getvalue(),
            file_name="计划数据表.xlsx",
            mime="application/vnd.ms-excel",
            key='download-excel',
            use_container_width=True
        )
